# -*- coding: utf-8 -*-
# created on 2016/11/30
import copy
from mathsolver.functions.shulie.ShuLieNumberSums import ShuLieGetNumSums
from mathsolver.functions.shulie.QianNXiangHe import ShuLieQiuHe
from mathsolver.functions.base import *
from mathsolver.functions.shulie.ShuLie_getSeqPoly import ShuLieGetSeqPoly
from mathsolver.functions.shulie.ShuLie_getTongXiang import ShuLieGetTongXiang


class ShuLieGetSums001(BaseFunction):
    """
    eg. 等差数列{a_{n}}的公差为2,若a_{2},a_{4},a_{8}成等比数列,则{a_{n}}的前n项和S_{n}=()
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], BaseSequence)
        seq = args[0].sympify()
        num = args[1].sympify()
        stepsolver = ShuLieGetTongXiang(self.known).solver(args[0])
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        tongxiangs = stepsolver.output[0].sympify()
        seq_values = tongxiangs[seq]
        n_sums = []
        for seq_value in seq_values:
            stepsolver = ShuLieQiuHe().solver([seq, seq_value], 1, num)
            n_sums_values = stepsolver.output[0].sympify()
            self.steps += stepsolver.steps
            n_sums.extend(n_sums_values)
        self.output.append(BaseNumbers(n_sums))
        return self


class ShuLieGetSums002(BaseFunction):
    """
    eg. 已知等比数列的公比为2,且前5项和为1,则前10项的和等于(   )
    """
    def solver(self, *args):
        assert len(args) == 1
        sls = [self.known[item] for item in self.known if isinstance(self.known[item], BaseSequence)]
        assert len(sls) == 1
        sl = sls[0]
        seq = sl.sympify()
        num = args[0].sympify()
        stepsolver = ShuLieGetTongXiang(self.known).solver(sl)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        tongxiangs = stepsolver.output[0].sympify()
        seq_values = tongxiangs[seq]
        n_sums = []
        for seq_value in seq_values:
            stepsolver = ShuLieQiuHe().solver([seq, seq_value], 1, num)
            n_sums_values = stepsolver.output[0].sympify()
            self.steps += stepsolver.steps
            n_sums.extend(n_sums_values)
        self.output.append(BaseNumbers(n_sums))
        return self


class ShuLieGetSums003(BaseFunction):
    """
    eg.已知等差数列{a_{n}}的前n项和为S_{n},a_{5}=5,S_{5}=15,则数列{\\frac{1}{{a}_{n}{a}_{n+1}}}的前100项和为()
    """
    def solver(self, *args):
        assert len(args) == 2
        assert isinstance(args[0], (BaseSeqPoly, BasePoly))
        seq = args[0].sympify()
        num = args[1].sympify()
        stepsolver = ShuLieGetSeqPoly(self.known).solver(args[0])
        self.steps += stepsolver.steps
        seq_values = stepsolver.output[0].sympify()
        n_sums = []
        seq_name = sympify("t_(n)")
        self.steps.append(["", "设%s = %s, 则" % (new_latex(seq_name), new_latex(seq))])
        for seq_value in seq_values:
            stepsolver = ShuLieQiuHe().solver([seq_name, seq_value], 1, num)
            n_sums_values = stepsolver.output[0].sympify()
            self.steps += stepsolver.steps
            n_sums.extend(n_sums_values)
        self.output.append(BaseNumbers(n_sums))
        return self


class ShuLieGetSums004(BaseFunction):
    """
    数列{a_{n}}满足a_{n+1}+(-1)^{n}a_{n}=2n-1,则{a_{n}}的前60项和为()
    """
    def solver(self, *args):
        stepsolver = ShuLieGetNumSums(known=self.known).solver(*args)
        self.steps += stepsolver.steps
        self.label.update(stepsolver.label)
        return self


class ShuLieGetSums(BaseFunction):
    CLS = [ShuLieGetSums001, ShuLieGetSums002, ShuLieGetSums003]

    def solver(self, *args):
        r = None
        for cl in ShuLieGetSums.CLS:
            try:
                new_known = copy.deepcopy(self.known)
                r = cl(new_known, verbose=True).solver(*args)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r
